package com.example.demo.ai.train;

import com.example.demo.ai.sftp.SftpService;
import com.jcraft.jsch.*;
import org.json.JSONException;
import org.json.JSONObject;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;

@Service
public class TrainService {

    @Autowired
    private SftpService sftpService;

    public String executeRemoteCommand(String yamlPath, String pyPath) throws Exception {

        Channel channel = sftpService.createChannel("exec");
        ChannelExec exec = (ChannelExec) channel;
        if(exec == null) return null;

        String command = String.format(
                "cd /root/yolo-project && " +
                        "source /root/yolo11/bin/activate && " +
                        "python3 /root/yolo-project/py/train.py --data %s --model %s",
                yamlPath, pyPath
        );

        exec.setCommand(command);

        // 解析Python返回的JSON结果
        String output = readCommandOutput(exec.getInputStream()); // 自定义方法读取输出流
        try {
            JSONObject result = new JSONObject(output);
            if ("success".equals(result.getString("status"))) {
                return result.getString("model_path");
            } else {
                throw new RuntimeException("训练失败: " + result.getString("message"));
            }
        } catch (JSONException e) {
            e.printStackTrace();
            return null;
        }

    }

    private String readCommandOutput(InputStream inputStream) throws IOException {
        // 使用StringBuilder高效拼接大文本
        StringBuilder outputBuilder = new StringBuilder();

        // 使用try-with-resources自动关闭流
        try (BufferedReader reader = new BufferedReader(
                new InputStreamReader(inputStream, StandardCharsets.UTF_8))) {

            String line;
            // 逐行读取输出
            while ((line = reader.readLine()) != null) {
                outputBuilder.append(line).append("\n");
            }
        }

        return outputBuilder.toString();
    }
}
